package com.xiexie.gateway.config;

import com.alibaba.csp.sentinel.adapter.gateway.common.rule.GatewayFlowRule;
import com.alibaba.csp.sentinel.adapter.gateway.common.rule.GatewayRuleManager;
import com.alibaba.csp.sentinel.adapter.gateway.sc.callback.BlockRequestHandler;
import com.alibaba.csp.sentinel.slots.block.degrade.DegradeException;
import com.alibaba.csp.sentinel.slots.block.flow.FlowException;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.server.ServerResponse;

import javax.annotation.PostConstruct;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

/**
 * @Description 自定义网关限流降级和异常
 * @Date 2022-04-20 9:19
 * @Author xie
 */
@Configuration
public class GatewaySentinelConfig {

    @PostConstruct
    public void init() {

        // 自定义异常信息
        initBlockRequestHandler();

        // 自定义网关限流规则
        initGatewayRules();
    }

    private void initBlockRequestHandler() {
        BlockRequestHandler blockRequestHandler = (serverWebExchange, throwable) -> {

            // 自定义异常处理
            Map<String, String> result = new HashMap<>();

            if (throwable instanceof FlowException) {
                // 限流

            } else if (throwable instanceof DegradeException) {
                // 降级

            } // .......

            result.put("code", String.valueOf(HttpStatus.TOO_MANY_REQUESTS.value()));
            result.put("code", HttpStatus.TOO_MANY_REQUESTS.getReasonPhrase());

            return ServerResponse.status(HttpStatus.OK)
                    .contentType(MediaType.APPLICATION_JSON)
                    .body(BodyInserters.fromValue(result));
        };
    }

    private void initGatewayRules() {

        Set<GatewayFlowRule> rules = new HashSet<>();
        //resource：资源名称，可以是网关中的 route 名称或者用户自定义的 API 分组名称。
        //count：限流阈值
        //intervalSec：统计时间窗口，单位是秒，默认是 1 秒。
        rules.add(new GatewayFlowRule("order_route").setCount(2).setIntervalSec(1));
        rules.add(new GatewayFlowRule("user_server_api").setCount(2).setIntervalSec(1));

        // 加载网关规则
        GatewayRuleManager.loadRules(rules);
    }
}
